{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "0512.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/yananma/5_programs_per_day/blob/master/0512.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fdxGFd9MF1Iw",
        "colab_type": "text"
      },
      "source": [
        "## 10.3 word2vec 的实现"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sbvbn6gb3czm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install mxnet d2lzh"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DxhTRHuV4SRW",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 151
        },
        "outputId": "3e828c48-7257-44b1-b82d-80e528956513"
      },
      "source": [
        "!git clone https://www.github.com/d2l-ai/d2l-zh.git"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'd2l-zh'...\n",
            "warning: redirecting to https://github.com/d2l-ai/d2l-zh.git/\n",
            "remote: Enumerating objects: 17, done.\u001b[K\n",
            "remote: Counting objects:   5% (1/17)\u001b[K\rremote: Counting objects:  11% (2/17)\u001b[K\rremote: Counting objects:  17% (3/17)\u001b[K\rremote: Counting objects:  23% (4/17)\u001b[K\rremote: Counting objects:  29% (5/17)\u001b[K\rremote: Counting objects:  35% (6/17)\u001b[K\rremote: Counting objects:  41% (7/17)\u001b[K\rremote: Counting objects:  47% (8/17)\u001b[K\rremote: Counting objects:  52% (9/17)\u001b[K\rremote: Counting objects:  58% (10/17)\u001b[K\rremote: Counting objects:  64% (11/17)\u001b[K\rremote: Counting objects:  70% (12/17)\u001b[K\rremote: Counting objects:  76% (13/17)\u001b[K\rremote: Counting objects:  82% (14/17)\u001b[K\rremote: Counting objects:  88% (15/17)\u001b[K\rremote: Counting objects:  94% (16/17)\u001b[K\rremote: Counting objects: 100% (17/17)\u001b[K\rremote: Counting objects: 100% (17/17), done.\u001b[K\n",
            "remote: Compressing objects: 100% (13/13), done.\u001b[K\n",
            "remote: Total 15702 (delta 9), reused 8 (delta 4), pack-reused 15685\u001b[K\n",
            "Receiving objects: 100% (15702/15702), 159.56 MiB | 34.00 MiB/s, done.\n",
            "Resolving deltas: 100% (11132/11132), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G1n-g3z24RGO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir ../data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y_KsyDQI4QVs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!cp ./d2l-zh/data/ptb.zip ../data/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iM_y6BJw3ib6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import collections \n",
        "import d2lzh as d2l \n",
        "import math \n",
        "from mxnet import autograd, gluon, nd \n",
        "from mxnet.gluon import data as gdata, loss as gloss, nn \n",
        "import random \n",
        "import sys \n",
        "import time \n",
        "import zipfile "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1bXdskL84KSE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "9ff113c1-760c-4c13-c12c-9426a21f2fc1"
      },
      "source": [
        "with zipfile.ZipFile('../data/ptb.zip', 'r') as zin:\n",
        "    zin.extractall('../data/')\n",
        "\n",
        "with open('../data/ptb/ptb.train.txt', 'r') as f:\n",
        "    lines = f.readlines()\n",
        "    raw_dataset = [st.split() for st in lines]\n",
        "\n",
        "'# sentences: %d' % len(raw_dataset)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'# sentences: 42068'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CPIu8vpL5oxO",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "outputId": "3343a9f7-662b-4535-e81a-d4108f240b16"
      },
      "source": [
        "for st in raw_dataset[:3]:\n",
        "    print('# tokens:', len(st), st[:5])"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "# tokens: 24 ['aer', 'banknote', 'berlitz', 'calloway', 'centrust']\n",
            "# tokens: 15 ['pierre', '<unk>', 'N', 'years', 'old']\n",
            "# tokens: 11 ['mr.', '<unk>', 'is', 'chairman', 'of']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o5w_JOjs6WK_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "counter = collections.Counter([tk for st in raw_dataset for tk in st])\n",
        "counter = dict(filter(lambda x: x[1] >= 5, counter.items()))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yZ6C2W9-85sh",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "60593bde-e104-4222-8071-9f2ecf7e3ce1"
      },
      "source": [
        "idx_to_token = [tk for tk, _ in counter.items()]\n",
        "token_to_idx = {tk: idx for idx, tk in enumerate(idx_to_token)}\n",
        "dataset = [[token_to_idx[tk] for tk in st if tk in token_to_idx] for st in raw_dataset]\n",
        "num_tokens = sum([len(st) for st in dataset])\n",
        "\"# tokens: %d\" % num_tokens"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'# tokens: 887100'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "884V0FQo9uQG",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "300bf8d1-d2f9-40b0-9c04-b3f4a680b314"
      },
      "source": [
        "def discard(idx):\n",
        "    return random.uniform(0, 1) < 1 - math.sqrt(1e-4 / counter[idx_to_token[idx]] * num_tokens)\n",
        "\n",
        "subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset]\n",
        "'# tokens: %d' % sum([len(st) for st in subsampled_dataset])"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'# tokens: 375509'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YOmEi6lJ_CBe",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "26677de0-c672-48a1-f01c-81a2fc87ab8f"
      },
      "source": [
        "def compare_counts(token):\n",
        "    return '#%s: before=%d, after=%d' %(token, \n",
        "        sum([st.count(token_to_idx[token]) for st in dataset]), \n",
        "        sum([st.count(token_to_idx[token]) for st in subsampled_dataset]))\n",
        "\n",
        "compare_counts('the')"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'#the: before=50770, after=2075'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n4W4-Jr2AMN_",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "58442f32-d3cf-49ed-ebff-d498743d6448"
      },
      "source": [
        "compare_counts('join')"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'#join: before=45, after=45'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4YEY0zCrBu2l",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_centers_and_contexts(dataset, max_window_size):\n",
        "    centers, contexts = [], []\n",
        "    for st in dataset:\n",
        "        if len(st) < 2:\n",
        "            continue \n",
        "        centers += st \n",
        "        for center_i in range(len(st)):\n",
        "            window_size = random.randint(1, max_window_size)\n",
        "            indices = list(range(max(0, center_i - window_size), min(len(st), center_i + 1 + window_size)))\n",
        "            indices.remove(center_i)\n",
        "            contexts.append([st[idx] for idx in indices])\n",
        "    return centers, contexts"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hkX3keYIDA7Q",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 202
        },
        "outputId": "3fb63aa0-9ffe-453d-9224-5a110210a3d1"
      },
      "source": [
        "tiny_dataset = [list(range(7)), list(range(7, 10))]\n",
        "print('dataset', tiny_dataset)\n",
        "for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):\n",
        "    print('center', center, 'has contexts', context)"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "dataset [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]]\n",
            "center 0 has contexts [1, 2]\n",
            "center 1 has contexts [0, 2]\n",
            "center 2 has contexts [0, 1, 3, 4]\n",
            "center 3 has contexts [2, 4]\n",
            "center 4 has contexts [2, 3, 5, 6]\n",
            "center 5 has contexts [3, 4, 6]\n",
            "center 6 has contexts [5]\n",
            "center 7 has contexts [8, 9]\n",
            "center 8 has contexts [7, 9]\n",
            "center 9 has contexts [8]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JNDaOk2_DkZ6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MXkGlFt-D1no",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_negatives(all_context, sampling_weights, K):\n",
        "    all_negatives, neg_candidates, i = [], [], 0 \n",
        "    population = list(range(len(sampling_weights)))\n",
        "    for contexts in all_contexts:\n",
        "        negatives = []\n",
        "        while len(negatives) < len(contexts) * K:\n",
        "            if i == len(neg_candidates):\n",
        "                i, neg_candidates = 0, random.choices(population, sampling_weights, k=int(1e5))\n",
        "            neg, i = neg_candidates[i], i + 1 \n",
        "            if neg not in set(contexts):\n",
        "                negatives.append(neg)\n",
        "        all_negatives.append(negatives)\n",
        "    return all_negatives \n",
        "\n",
        "sampling_weights = [counter[w]**0.75 for w in idx_to_token]\n",
        "all_negatives = get_negatives(all_contexts, sampling_weights, 5)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lIWO5rLfGFp1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def batchify(data):\n",
        "    max_len = max(len(c) + len(n) for _, c, n in data)\n",
        "    centers, contexts_negatives, masks, labels = [], [], [], []\n",
        "    for center, context, negative in data:\n",
        "        cur_len = len(context) + len(negative)\n",
        "        centers += [center]\n",
        "        contexts_negatives += [context + negative + [0] * (max_len - cur_len)]\n",
        "        masks += [[1] * cur_len + [0] * (max_len - cur_len)]\n",
        "        labels += [[1] * len(context) + [0] * (max_len - len(context))]\n",
        "\n",
        "    return (nd.array(centers).reshape((-1, 1)), nd.array(contexts_negatives), \n",
        "        nd.array(masks), nd.array(labels))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PBRF3K0UwCR9",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 84
        },
        "outputId": "8334e16c-fb1d-46a8-cc01-452ef9fb0603"
      },
      "source": [
        "batch_size = 512 \n",
        "num_workers = 4 \n",
        "dataset = gdata.ArrayDataset(all_centers, all_contexts, all_negatives)\n",
        "data_iter = gdata.DataLoader(dataset, batch_size, shuffle=True, batchify_fn=batchify, num_workers=num_workers)\n",
        "for batch in data_iter:\n",
        "    for name, data in zip(['centers', 'contexts_negatives', 'masks', 'labels'], batch):\n",
        "        print(name, 'shape:', data.shape)\n",
        "    break "
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "centers shape: (512, 1)\n",
            "contexts_negatives shape: (512, 60)\n",
            "masks shape: (512, 60)\n",
            "labels shape: (512, 60)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lrSRzjNHxa-B",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "05bf5004-708c-43c3-829d-4d0187b2ca60"
      },
      "source": [
        "embed = nn.Embedding(input_dim=20, output_dim=4)\n",
        "embed.initialize()\n",
        "embed.weight"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Parameter embedding0_weight (shape=(20, 4), dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HAVGTLuM2iO7",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 168
        },
        "outputId": "0bd0e76e-d90b-457c-8750-60f62766f435"
      },
      "source": [
        "x = nd.array([[1, 2, 3], [4, 5, 6]])\n",
        "embed(x)"
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[[[ 0.01438687  0.05011239  0.00628365  0.04861524]\n",
              "  [-0.01068833  0.01729892  0.02042518 -0.01618656]\n",
              "  [-0.00873779 -0.02834515  0.05484822 -0.06206018]]\n",
              "\n",
              " [[ 0.06491279 -0.03182812 -0.01631819 -0.00312688]\n",
              "  [ 0.0408415   0.04370362  0.00404529 -0.0028032 ]\n",
              "  [ 0.00952624 -0.01501013  0.05958354  0.04705103]]]\n",
              "<NDArray 2x3x4 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IxogaRc92wod",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "860f3030-1fd0-41a4-9540-18e0b4ed1508"
      },
      "source": [
        "X = nd.ones((2, 1, 4))\n",
        "Y = nd.ones((2, 4, 6))\n",
        "nd.batch_dot(X, Y).shape "
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(2, 1, 6)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 38
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sd58yhTu3BpU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def skip_gram(center, contexts_and_negatives, embed_v, embed_u):\n",
        "    v = embed_v(center)\n",
        "    u = embed_u(contexts_and_negatives)\n",
        "    pred = nd.batch_dot(v, u.swapaxes(1, 2))\n",
        "    return pred "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "20T_bwPt3nMZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loss = gloss.SigmoidBinaryCrossEntropyLoss()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6GSVHWEW3zzS",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "outputId": "cc7b07b3-7bf7-460b-8614-48424cc35fce"
      },
      "source": [
        "pred = nd.array([[1.5, 0.3, -1, 2], [1.1, -0.6, 2.2, 0.4]])\n",
        "label = nd.array([[1, 0, 0, 0], [1, 1, 0, 0]])\n",
        "mask = nd.array([[1, 1, 1, 1], [1, 1, 1, 0]])\n",
        "loss(pred, label, mask) * mask.shape[1] / mask.sum(axis=1)"
      ],
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\n",
              "[0.8739896 1.2099689]\n",
              "<NDArray 2 @cpu(0)>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 57
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MF-4GiKV5N_f",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 50
        },
        "outputId": "49e0948b-16bd-44d2-8cc2-40ece1e2396e"
      },
      "source": [
        "def sigmd(x):\n",
        "    return -math.log(1 / (1 + math.exp(-x)))\n",
        "\n",
        "print('%.7f' % ((sigmd(1.5) + sigmd(-0.3) + sigmd(1) + sigmd(-2)) / 4))\n",
        "print('%.7f' % ((sigmd(1.1) + sigmd(-0.6) + sigmd(-2.2)) / 3))"
      ],
      "execution_count": 58,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "0.8739896\n",
            "1.2099689\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0lF1RigX7XKi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "embed_size = 100 \n",
        "net = nn.Sequential()\n",
        "net.add(nn.Embedding(input_dim=len(idx_to_token), output_dim=embed_size),\n",
        "    nn.Embedding(input_dim=len(idx_to_token), output_dim=embed_size))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I2fxNkh08CaI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(net, lr, num_epochs):\n",
        "    ctx = d2l.try_gpu()\n",
        "    net.initialize(ctx=ctx, force_reinit=True)\n",
        "    trainer = gluon.Trainer(net.collect_params(), 'adam', {'learning_rate': lr})\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        start, l_sum, n = time.time(), 0.0, 0 \n",
        "        for batch in data_iter:\n",
        "            center, context_negative, mask, label = [data.as_in_context(ctx) for data in batch]\n",
        "            with autograd.record():\n",
        "                pred = skip_gram(center, context_negative, net[0], net[1])\n",
        "                l = (loss(pred.reshape(label.shape), label, mask) * mask.shape[1] / mask.sum(axis=1))\n",
        "            l.backward()\n",
        "            trainer.step(batch_size)\n",
        "            l_sum += l.sum().asscalar()\n",
        "            n += l.size \n",
        "        print('epoch %d, loss %.2f, time %.2fs' % (epoch + 1, l_sum / n, time.time() - start))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yYn_PIrl908t",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 101
        },
        "outputId": "9003ec95-cf23-456d-b5a0-719a7054e2e3"
      },
      "source": [
        "train(net, 0.005, 5)"
      ],
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "epoch 1, loss 0.46, time 158.52s\n",
            "epoch 2, loss 0.39, time 155.11s\n",
            "epoch 3, loss 0.35, time 147.60s\n",
            "epoch 4, loss 0.32, time 148.85s\n",
            "epoch 5, loss 0.30, time 154.30s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9FG_1LWt94Jc",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "outputId": "1dd20516-55e6-4fc1-d0cb-8d44d23aa573"
      },
      "source": [
        "def get_similar_tokens(query_token, k, embed):\n",
        "    W = embed.weight.data()\n",
        "    x = W[token_to_idx[query_token]]\n",
        "    cos = nd.dot(W, x) / (nd.sum(W * W, axis=1) * nd.sum(x * x) + 1e-9).sqrt()\n",
        "    topk = nd.topk(cos, k=k+1, ret_typ='indices').asnumpy().astype('int32')\n",
        "    for i in topk[1:]:\n",
        "        print('cosine sim=%.3f: %s' % (cos[i].asscalar(), (idx_to_token[1])))\n",
        "\n",
        "get_similar_tokens('chip', 3, net[0])"
      ],
      "execution_count": 67,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "cosine sim=0.516: <unk>\n",
            "cosine sim=0.480: <unk>\n",
            "cosine sim=0.470: <unk>\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}